#!/usr/bin/env python3
# coding=utf-8
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import numpy as  np
import queue as Q
import pickle
import os
from math import sin, cos, sqrt, atan2, radians
import argparse
import pandas as pd
from scipy import stats

def get_distance(lat_1, lng_1, lat_2, lng_2,DEBUG=0):
    # radius of earth in km
    R = 6373.0

    dlng = radians(lat_2) - radians(lat_1)
    dlat = radians(lng_2) - radians(lng_1)

    x = dlat*cos((lat_1+lat_2)*0.5)

    d = R * sqrt(x**2 + dlng**2)
    # if DEBUG!=0:
    #     print (d)
    return d


def kNN(sensor_id_list, lat_list, long_list, id, k=2,
        DEBUG=0):
    kNN_PQ = Q.PriorityQueue(maxsize=k)

    nearest_distance = np.inf
    nearest_id = -1
    for _id in sensor_id_list:
        if _id == id:
            continue
        dis = -get_distance(lat_list[id-1], long_list[id-1],
                            lat_list[_id-1], long_list[_id-1],
                            DEBUG)
        if not kNN_PQ.full():
            kNN_PQ.put([dis, _id])
        else:
            temp_data = kNN_PQ.get()
            temp_dis = temp_data[0]
            temp_label = temp_data[1]
            if dis > temp_dis:
                temp_dis = dis
                temp_label = _id
            kNN_PQ.put([temp_dis, temp_label])
    return kNN_PQ

def get_adjacency_matrix(sensor_id_list, sensor_name,
                         lat_list, long_list, k=2,
                         DEBUG=0):
    """
    :sensor_id_list: list of sensor id
    :lat_list:       list of latitude of sensor
    :long_list:      list of longtitude of sensor
    """
    sensor_num = len(sensor_id_list)
    dist_mx = np.zeros((sensor_num, sensor_num), dtype=np.float32)
    sensor_name2id = {}
    # Construct the graph using kNN, add edge to k nearest neighors
    for idx, id in enumerate(sensor_id_list):
        # if DEBUG!=0:
        #     print (sensor_name[idx], id-1)
        sensor_name2id[sensor_name[idx]] = id-1
        knn_pq = kNN(sensor_id_list, lat_list, long_list, id, k,
                     DEBUG)
        knn_list = knn_pq.queue
        for nei in knn_list:
            dist_mx[id - 1, nei[1] - 1] = 1
            dist_mx[id - 1][id - 1] = 1
    if DEBUG!=0:
        print("dist_mx=",stats.describe(dist_mx))
    return sensor_name2id, dist_mx

def generate_adj_matrix_pkl(input_file,output_dir,
                            id_col="0",long_col="1",lat_col="2",
                            DEBUG=0
                            ):
    if not os.path.exists(output_dir):
        print("[ERROR] Output dir not exists!")
    else:
        if isinstance(input_file, str) and  os.path.exists(input_file):

            print("[INFO] Loading input data {path} for computing adj mat".format(path=input_file))
            sensor_info = pd.read_csv(input_file)
            print("[INFO] Shape of Data {shape}".format(shape=sensor_info.shape) )
            sensor_id_list =list(sensor_info.index)
            sensor_name =[]
            lat_list = []
            long_list = []

            if id_col.isdigit():
                sensor_name = sensor_info[sensor_info.columns[int( id_col )]].tolist()
            else:
                sensor_name = sensor_info[id_col].tolist()

            if long_col.isdigit():
                long_list = sensor_info[sensor_info.columns[ int( long_col ) ]].tolist()
            else:
                long_list = sensor_info[long_col].tolist()

            if lat_col.isdigit():
                lat_list = sensor_info[sensor_info.columns[ int( lat_col ) ]].tolist()
            else:
                lat_list = sensor_info[lat_col].tolist()

            sensor_name2id, adj_mx = get_adjacency_matrix(sensor_id_list,
                                                        sensor_name,
                                                        lat_list,
                                                        long_list,
                                                          2,DEBUG)
            pkl_outputpath = os.path.join(output_dir,'adj_mat_ori.pkl')
            print("[INFO] Writing adj matrix to {path}".format(
                path=pkl_outputpath))
            if DEBUG:
                print(sensor_name)
                print(sensor_name2id)
            with open(pkl_outputpath , 'wb') as f:
                # pickle.dump([sensor_name, sensor_name2id, adj_mx], f)
                pickle.dump([sensor_name, sensor_name2id, adj_mx], f,
                            protocol=2)
        else:
            print("[ERROR] Input file path Error !")
            print("[INFO] Value is {value}...".format(value=input_file))

def main(args):
    print(args.DEBUG)
    generate_adj_matrix_pkl(input_file = args.input_file,
                            output_dir = args.output_dir,
                            id_col= args.id_col,
                            long_col= args.long_col,
                            lat_col= args.lat_col,
                            DEBUG=args.DEBUG)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output_dir" , type=str, default="data/METR-LA/", help="Output directory"
    )
    parser.add_argument(
        "--input_file", type=str, default="data/METR-LA/coordinates.csv",help="Coordinates of observations"
    )
    parser.add_argument(
        "--id_col", type=str,default="0",help="Column ID name or Index"
    )
    parser.add_argument(
        "--long_col", type=str,default="1",help="Column longitude name or Index"
    )
    parser.add_argument(
        "--lat_col", type=str,default="2",help="Column latitude name or Index"
    )
    parser.add_argument(
        "--DEBUG",type=int, default=0, help="Whether(1/0) to print Debug info"
    )
    args = parser.parse_args()
    main(args)
    # aq_station_table='/media/data/seifer_08ms/dl/gnn/traffic-prediction/repos/gc-dcrnn/data/kdd_cup/Beijing_AirQuality_Stations_cn.csv'
